package org.lanyu.springainovel.chat.service;

import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.lanyu.springainovel.common.config.AIModelInitService;
import org.lanyu.springainovel.common.config.ChatModelManager;
import org.lanyu.springainovel.common.config.VectorStoreManager;
import org.lanyu.springainovel.common.entity.AiModel;
import org.lanyu.springainovel.common.entity.MyPromptTemplate;
import org.lanyu.springainovel.common.entity.User;
import org.lanyu.springainovel.common.repository.CustomJdbcChatMemoryRepository;
import org.lanyu.springainovel.common.service.AiModelService;
import org.lanyu.springainovel.common.service.DynamicMcpClientManager;
import org.lanyu.springainovel.common.service.UserCacheService;
import org.lanyu.springainovel.chat.vo.QuestionVO;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.vectorstore.QuestionAnswerAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.content.Media;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.ByteArrayResource;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.context.SecurityContextHolder;
import org.springframework.stereotype.Service;
import org.springframework.util.MimeTypeUtils;
import org.springframework.web.reactive.function.client.WebClientRequestException;
import reactor.core.publisher.Flux;

import java.util.ArrayList;
import java.util.Base64;
import java.util.List;

/**
 * 聊天流式处理服务�? * 负责处理复杂的聊天逻辑，包括知识库检索、多模态处理、工具调用等
 */
@Slf4j
@Service
@AllArgsConstructor
public class ChatStreamService {

    private final ChatMemoryService chatMemoryService;
    private final ChatModelManager chatModelManager;
    private final AiModelService aiModelService;
    private final DynamicMcpClientManager dynamicMcpClientManager;
    private final AIModelInitService aiModelInitService;
    private final UserCacheService userCacheService;
    private final CustomJdbcChatMemoryRepository customJdbcChatMemoryRepository;
    private final PromptTemplateService promptTemplateService;
    private final VectorStoreManager vectorStoreManager;

    /**
     * 处理流式聊天请求
     *
     * @param body 用户请求参数
     * @return 流式响应
     */
    public Flux<String> processStreamChat(QuestionVO body) {
        try {
            // 获取当前用户（如果已认证�?
            User currentUser = getCurrentUser();
            final User finalCurrentUser = currentUser;

            // 验证会话ID
            if (body.getSessionId() == null || body.getSessionId().isEmpty()) {
                return Flux.just("错误: 会话ID不能为空，请确保前端传递了有效的会话ID");
            }

            final String finalSessionId = body.getSessionId();

            // 更新会话和对话记�?
            if (finalCurrentUser != null && body.getMemory() && body.getContentIndex() != null) {
                chatMemoryService.saveOrUpdateSessionHistory(body);
            }

            // 获取聊天模型
            AiModel modelConfig = getModelConfig(body.getModelId());
            if (modelConfig == null) {
                return Flux.just("错误: 未找到指定模型ID: " + body.getModelId());
            }

            ChatModel chatModel;
            try {
                chatModel = chatModelManager.getOrCreateChatModel(modelConfig);
            } catch (Exception e) {
                // 捕获模型创建异常，可能是模型不在线或服务不可�?
                log.error("创建聊天模型失败: {}", e.getMessage());
                if (e.getCause() instanceof WebClientRequestException ||
                        e.getMessage().contains("Connection refused") ||
                        e.getMessage().contains("ConnectException")) {
                    return Flux.just("该模型暂不支持");
                }
                return Flux.just("模型创建失败: " + e.getMessage());
            }

            // 创建ChatClient并添加MessageChatMemoryAdvisor（如果启用了记忆功能�?
            ChatClient chatClient = createChatClient(chatModel, body);

            // 生成流式响应
            Flux<String> response = generateStreamResponse(chatClient, body, finalSessionId);

            return response
                    .doOnComplete(() -> {
                        log.info("会话 {} 完成，使用知识库: {}", finalSessionId, body.getKnowledgeId() != null);
                    })
                    .onErrorResume(throwable -> {
                        // 捕获所有异常并转换为Flux<String>
                        log.error("流式接口发生异常", throwable);

                        // 特殊处理模型连接异常
                        if (throwable instanceof WebClientRequestException ||
                                (throwable.getCause() instanceof WebClientRequestException)) {
                            log.warn("模型连接失败，可能是模型不在线或服务不可用");
                            return Flux.just("模型不在线或服务不可用");
                        }

                        // 其他异常返回通用错误信息
                        return Flux.just("错误: " + throwable.getMessage());
                    });
        } catch (Exception e) {
            // 捕获所有同步异常并转换为Flux<String>
            log.error("流式接口同步异常", e);
            return Flux.just("系统错误: " + e.getMessage());
        }
    }

    /**
     * 获取当前用户
     */
    private User getCurrentUser() {
        Authentication authentication = SecurityContextHolder.getContext().getAuthentication();
        if (authentication != null && authentication.isAuthenticated() &&
                !authentication.getName().equals("anonymousUser")) {
            return userCacheService.getUserByUsername(authentication.getName());
        }
        return null;
    }

    /**
     * 获取模型配置
     */
    private AiModel getModelConfig(Long modelId) {
        return aiModelService.findById(modelId);
    }

    /**
     * 创建ChatClient
     */
    private ChatClient createChatClient(ChatModel chatModel, QuestionVO body) {
        // 创建ChatClient并添加MessageChatMemoryAdvisor（如果启用了记忆功能）
        ChatClient.Builder chatClientBuilder = ChatClient.builder(chatModel);

        // 如果启用了记忆功能，添加MessageChatMemoryAdvisor
        if (body.getMemory()) {
            // 使用自定义的数据库存储来创建MessageWindowChatMemory
            ChatMemory chatMemory = customJdbcChatMemoryRepository.getChatMemory(body.getSessionId(), body.getMaxTokens());

            chatClientBuilder.defaultAdvisors(MessageChatMemoryAdvisor.builder(chatMemory)
                    .build());
        }

        // 添加知识库检索.检查向量存储是否已启用
        if (body.getKnowledgeId() != null && vectorStoreManager.isVectorStoreEnabled()) {
            User user = getCurrentUser();
            String indexName = "kb_" + body.getKnowledgeId();

            // 检查向量存储是否已启用
            VectorStore vectorStore = aiModelInitService.getVectorStoreByUser(user, indexName);

            if (vectorStore != null) {
                // 创建提示词模板
                PromptTemplate promptTemplate = new PromptTemplate(
                        "请基于以下参考内容回答用户问题，如果参考内容中没有相关信息，请诚实地表示不知道。\n\n【参考内容】：\n{documents}\n\n【用户问题】：\n{question}");

                // 使用QuestionAnswerAdvisor的Builder模式创建，并设置提示词模板
                chatClientBuilder.defaultAdvisors(
                        QuestionAnswerAdvisor.builder(vectorStore)
                                .searchRequest(SearchRequest.builder()
                                        .similarityThreshold(0.6)
                                        .topK(6)
                                        .build())
                                .promptTemplate(promptTemplate)
                                .build()
                );
            }
        }

        return chatClientBuilder.build();
    }

    /**
     * 生成流式响应
     */
    private Flux<String> generateStreamResponse(ChatClient chatClient, QuestionVO body, String sessionId) {
        // 创建ChatClient的prompt spec
        var promptSpec = chatClient.prompt();

        // 如果启用了记忆功能，设置会话ID
        if (body.getMemory()) {
            promptSpec = promptSpec.advisors(advisorSpec -> advisorSpec
                    .param(ChatMemory.CONVERSATION_ID, sessionId));
        }

        // 处理系统提示�?
        if (body.getSystemPromptId() != null && body.getSystemPromptId() > 0) {
            try {
                MyPromptTemplate template = promptTemplateService.getTemplateById(body.getSystemPromptId());
                if (template != null && template.getContent() != null && !template.getContent().isEmpty()) {
                    promptSpec = promptSpec.system(template.getContent());
                    // 增加模板使用次数
                    promptTemplateService.incrementUsageCount(body.getSystemPromptId());
                }
            } catch (Exception e) {
                log.error("获取系统提示词模板失�? {}", e.getMessage());
            }
        }

        // 判断是否多模�?
        if (body.getFilebase64s() != null && !body.getFilebase64s().isEmpty()) {
            // 构建媒体列表
            List<Media> mediaList = new ArrayList<>();
            for (String fileBase64 : body.getFilebase64s()) {
                byte[] imageBytes = Base64.getDecoder().decode(fileBase64);
                ByteArrayResource imageResource = new ByteArrayResource(imageBytes);
                mediaList.add(new Media(MimeTypeUtils.IMAGE_PNG, imageResource));
            }

            // 使用正确的多模态方式调�?
            if (body.getUseTools() != null && body.getUseTools()) {
                List<ToolCallback> toolCallbacks = dynamicMcpClientManager.getAvailableToolCallbacks(true);
                return promptSpec.user(u -> u.text(body.getMessage())
                                .media(mediaList.toArray(new Media[0])))
                        .toolCallbacks(toolCallbacks).stream().content();
            } else {
                return promptSpec.user(u -> u.text(body.getMessage())
                        .media(mediaList.toArray(new Media[0]))).stream().content();
            }
        } else {
            // 使用纯文本方式调�?
            if (body.getUseTools() != null && body.getUseTools()) {
                List<ToolCallback> toolCallbacks = dynamicMcpClientManager.getAvailableToolCallbacks(true);
                return promptSpec.user(body.getMessage()).toolCallbacks(toolCallbacks).stream().content();
            } else {
                return promptSpec.user(body.getMessage()).stream().content();
            }
        }
    }
}
